#!/usr/bin/env python3
# coding:utf-8
"""
Navigation Manager for MutoLib
Provides navigation-related functionalities:
1. Check navigation node status
2. Navigate to specified coordinates
3. Start/stop navigation system
4. Manage laser components
"""

import os
import sys
import time
import subprocess
import threading
from typing import Optional, Dict, List, Tuple

try:
    import rclpy
    from rclpy.node import Node
    from rclpy.action import ActionClient
    from rclpy.executors import MultiThreadedExecutor
    from geometry_msgs.msg import PoseStamped, PoseWithCovarianceStamped
    from nav2_msgs.action import NavigateToPose
    from lifecycle_msgs.msg import State
    from lifecycle_msgs.srv import GetState
    from std_msgs.msg import Header
    from builtin_interfaces.msg import Time
    from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
    ROS2_AVAILABLE = True
except ImportError:
    ROS2_AVAILABLE = False
    print("Warning: ROS2 packages not available. Navigation functionality will be limited.")


class NavigationManager:
    """
    Navigation Manager class providing navigation functionalities
    """
    
    def __init__(self, muto_instance=None):
        """
        Initialize NavigationManager
        
        Args:
            muto_instance: Reference to main MutoLargemodelInterface instance
        """
        self.muto = muto_instance
        self._ros_node = None
        self._nav_to_pose_client = None
        self._initial_pose_pub = None
        self._lifecycle_clients = {}
        self._launch_processes = {}
        
        # Navigation status
        self._navigation_result = None
        self._navigation_feedback = None
        self._is_navigating = False
        
        # Initialize ROS2 if available
        if ROS2_AVAILABLE:
            self._initialize_ros2()
    
    def _initialize_ros2(self):
        """
        Initialize ROS2 components
        """
        try:
            if not rclpy.ok():
                rclpy.init()
            
            self._ros_node = rclpy.create_node('navigation_manager')
            
            # Navigation action client
            self._nav_to_pose_client = ActionClient(
                self._ros_node, NavigateToPose, 'navigate_to_pose'
            )
            
            # Publisher for initial pose
            qos_profile = QoSProfile(
                reliability=ReliabilityPolicy.RELIABLE,
                history=HistoryPolicy.KEEP_LAST,
                depth=1
            )
            self._initial_pose_pub = self._ros_node.create_publisher(
                PoseWithCovarianceStamped, 
                'initialpose', 
                qos_profile
            )
            
            print("✅ Navigation Manager ROS2 components initialized")
            
        except Exception as e:
            print(f"❌ Failed to initialize ROS2 components: {e}")
            self._ros_node = None
    
    def _create_lifecycle_client(self, node_name: str):
        """
        Create lifecycle service client for a node
        
        Args:
            node_name: Name of the node
            
        Returns:
            Service client or None if ROS2 not available
        """
        if not self._ros_node:
            return None
            
        service_name = f'/{node_name}/get_state'
        if service_name not in self._lifecycle_clients:
            self._lifecycle_clients[service_name] = self._ros_node.create_client(
                GetState, service_name
            )
        return self._lifecycle_clients[service_name]
    
    def check_node_status(self, node_name: str, timeout: float = 5.0) -> tuple:
        """
        Check the status of a navigation-related node
        
        Args:
            node_name: Name of the node to check
            timeout: Timeout for service call
            
        Returns:
            Tuple (success, message, data)
        """
        if not ROS2_AVAILABLE or not self._ros_node:
            return False, "ROS2 not available", {}
        
        try:
            client = self._create_lifecycle_client(node_name)
            
            if not client or not client.wait_for_service(timeout_sec=timeout):
                return False, f"Service /{node_name}/get_state not available", {
                    'node_name': node_name,
                    'status': 'unavailable',
                    'state': None
                }
            
            request = GetState.Request()
            future = client.call_async(request)
            
            # Wait for response
            start_time = time.time()
            while not future.done() and (time.time() - start_time) < timeout:
                rclpy.spin_once(self._ros_node, timeout_sec=0.1)
            
            if future.done():
                response = future.result()
                state_labels = {
                    State.PRIMARY_STATE_UNKNOWN: 'unknown',
                    State.PRIMARY_STATE_UNCONFIGURED: 'unconfigured',
                    State.PRIMARY_STATE_INACTIVE: 'inactive',
                    State.PRIMARY_STATE_ACTIVE: 'active',
                    State.PRIMARY_STATE_FINALIZED: 'finalized'
                }
                
                return True, "Node status retrieved successfully", {
                    'node_name': node_name,
                    'status': 'available',
                    'state': state_labels.get(response.current_state.id, 'unknown'),
                    'state_id': response.current_state.id
                }
            else:
                return False, f"Timeout waiting for response from {node_name}", {
                    'node_name': node_name,
                    'status': 'timeout',
                    'state': None
                }
                
        except Exception as e:
            return False, f"Error checking node status for {node_name}: {str(e)}", {
                'node_name': node_name,
                'status': 'error',
                'state': None
            }
    
    def check_navigation_system_status(self) -> tuple:
        """
        Check the status of all navigation-related nodes
        
        Returns:
            Tuple (success, message, data)
        """
        if not ROS2_AVAILABLE or not self._ros_node:
            return False, "ROS2 not available", {}
        
        nav_nodes = [
            'controller_server',
            'planner_server', 
            'behavior_server',
            'bt_navigator',
            'waypoint_follower',
            'map_server',
            'amcl'
        ]
        
        results = {}
        overall_status = 'healthy'
        
        for node in nav_nodes:
            success, message, data = self.check_node_status(node)
            results[node] = data
            
            if not success or data.get('status') != 'available' or data.get('state') != 'active':
                overall_status = 'unhealthy'
        
        return True, "Navigation system status checked", {
            'overall_status': overall_status,
            'nodes': results,
            'timestamp': time.time()
        }
    
    def set_initial_pose(self, x: float, y: float, yaw: float, frame_id: str = 'map') -> tuple:
        """
        Set the initial pose of the robot
        
        Args:
            x: X coordinate
            y: Y coordinate  
            yaw: Yaw angle in radians
            frame_id: Frame ID (default: 'map')
            
        Returns:
            Tuple (success, message, data)
        """
        if not ROS2_AVAILABLE or not self._ros_node or not self._initial_pose_pub:
            return False, "ROS2 or initial pose publisher not available", {}
        
        try:
            pose_msg = PoseWithCovarianceStamped()
            pose_msg.header.frame_id = frame_id
            pose_msg.header.stamp = self._ros_node.get_clock().now().to_msg()
            
            # Set position
            pose_msg.pose.pose.position.x = x
            pose_msg.pose.pose.position.y = y
            pose_msg.pose.pose.position.z = 0.0
            
            # Set orientation (convert yaw to quaternion)
            import math
            pose_msg.pose.pose.orientation.x = 0.0
            pose_msg.pose.pose.orientation.y = 0.0
            pose_msg.pose.pose.orientation.z = math.sin(yaw / 2.0)
            pose_msg.pose.pose.orientation.w = math.cos(yaw / 2.0)
            
            # Set covariance (6x6 matrix, but we only set the relevant elements)
            pose_msg.pose.covariance = [0.0] * 36
            pose_msg.pose.covariance[0] = 0.25  # x variance
            pose_msg.pose.covariance[7] = 0.25  # y variance
            pose_msg.pose.covariance[35] = 0.06853891945200942  # yaw variance
            
            self._initial_pose_pub.publish(pose_msg)
            
            return True, f"Initial pose set: x={x}, y={y}, yaw={yaw}", {
                'x': x,
                'y': y,
                'yaw': yaw,
                'frame_id': frame_id
            }
            
        except Exception as e:
            return False, f"Error setting initial pose: {str(e)}", {}
    
    def navigate_to_pose(self, x: float, y: float, yaw: float, 
                        frame_id: str = 'map', timeout: float = 300.0) -> tuple:
        """
        Navigate to a specified pose
        
        Args:
            x: Target X coordinate
            y: Target Y coordinate
            yaw: Target yaw angle in radians
            frame_id: Frame ID (default: 'map')
            timeout: Navigation timeout in seconds
            
        Returns:
            Tuple (success, message, data)
        """
        if not ROS2_AVAILABLE or not self._ros_node or not self._nav_to_pose_client:
            return False, "ROS2 or navigation client not available", {}
        
        if self._is_navigating:
            return False, "Navigation already in progress", {}
        
        try:
            # Wait for action server
            if not self._nav_to_pose_client.wait_for_server(timeout_sec=10.0):
                return False, "Navigation action server not available", {}
            
            # Create goal message
            goal_msg = NavigateToPose.Goal()
            goal_msg.pose.header.frame_id = frame_id
            goal_msg.pose.header.stamp = self._ros_node.get_clock().now().to_msg()
            
            # Set target position
            goal_msg.pose.pose.position.x = x
            goal_msg.pose.pose.position.y = y
            goal_msg.pose.pose.position.z = 0.0
            
            # Set target orientation (convert yaw to quaternion)
            import math
            goal_msg.pose.pose.orientation.x = 0.0
            goal_msg.pose.pose.orientation.y = 0.0
            goal_msg.pose.pose.orientation.z = math.sin(yaw / 2.0)
            goal_msg.pose.pose.orientation.w = math.cos(yaw / 2.0)
            
            print(f"🎯 Sending navigation goal: x={x}, y={y}, yaw={yaw}")
            
            # Send goal
            self._is_navigating = True
            self._navigation_result = None
            
            send_goal_future = self._nav_to_pose_client.send_goal_async(
                goal_msg,
                feedback_callback=self._navigation_feedback_callback
            )
            
            # Wait for goal acceptance
            start_time = time.time()
            while not send_goal_future.done() and (time.time() - start_time) < 10.0:
                rclpy.spin_once(self._ros_node, timeout_sec=0.1)
            
            if not send_goal_future.done():
                self._is_navigating = False
                return False, "Timeout waiting for goal acceptance", {}
            
            goal_handle = send_goal_future.result()
            if not goal_handle.accepted:
                self._is_navigating = False
                return False, "Navigation goal rejected", {}
            
            print("✅ Navigation goal accepted")
            
            # Wait for result
            result_future = goal_handle.get_result_async()
            start_time = time.time()
            
            while not result_future.done() and (time.time() - start_time) < timeout:
                rclpy.spin_once(self._ros_node, timeout_sec=0.1)
            
            self._is_navigating = False
            
            if result_future.done():
                result = result_future.result()
                success = result.status == 4  # SUCCEEDED
                
                return success, "Navigation completed" if success else "Navigation failed", {
                    'success': success,
                    'status_code': result.status,
                    'target_x': x,
                    'target_y': y,
                    'target_yaw': yaw
                }
            else:
                # Cancel navigation on timeout
                goal_handle.cancel_goal_async()
                return False, f"Navigation timeout after {timeout} seconds", {
                    'target_x': x,
                    'target_y': y,
                    'target_yaw': yaw
                }
                
        except Exception as e:
            self._is_navigating = False
            return False, f"Error during navigation: {str(e)}", {}
    
    def _navigation_feedback_callback(self, feedback_msg):
        """
        Callback for navigation feedback
        
        Args:
            feedback_msg: Feedback message from navigation
        """
        self._navigation_feedback = feedback_msg.feedback
        # Log current position occasionally
        if hasattr(feedback_msg.feedback, 'current_pose'):
            pose = feedback_msg.feedback.current_pose.pose
            print(f"🚀 Navigation progress: x={pose.position.x:.2f}, y={pose.position.y:.2f}")
    
    def cancel_navigation(self) -> tuple:
        """
        Cancel current navigation
        
        Returns:
            Tuple (success, message, data)
        """
        if not self._is_navigating:
            return False, "No navigation in progress", {}
        
        # Reset the flag
        self._is_navigating = False
        
        return True, "Navigation cancelled", {}
    
    def start_navigation_system(self, algorithm: str = 'dwa', 
                              map_path: Optional[str] = None, wait_time_s: float = 15.0) -> tuple:
        """
        Start the navigation system with specified algorithm
        
        Args:
            algorithm: Navigation algorithm ('dwa' or 'teb')
            map_path: Optional path to map file
            wait_time_s: Time to wait for initialization (default: 15.0 seconds)
            
        Returns:
            Tuple (success, message, data)
        """
        if algorithm not in ['dwa', 'teb']:
            return False, f"Invalid algorithm: {algorithm}. Use 'dwa' or 'teb'", {}
        
        try:
            # Step 1: Start laser bringup if needed
            laser_bringup_key = 'laser_bringup'
            if laser_bringup_key not in self._launch_processes or self._launch_processes[laser_bringup_key].poll() is not None:
                success, message, data = self.start_laser_bringup()
                if not success:
                    return False, f"Failed to start laser bringup: {message}", {}
            
            # Step 2: Start laser odometry if needed
            laser_odom_key = 'laser_odometry'
            if laser_odom_key not in self._launch_processes or self._launch_processes[laser_odom_key].poll() is not None:
                success, message, data = self.start_laser_odometry()
                if not success:
                    return False, f"Failed to start laser odometry: {message}", {}
            
            # Build navigation launch command
            launch_file = f'navigation_{algorithm}_launch.py'
            cmd = ['ros2', 'launch', 'yahboomcar_nav', launch_file]
            
            if map_path:
                if not os.path.exists(map_path):
                    return False, f"Map file not found: {map_path}", {}
                cmd.extend(['map:=' + map_path])
            
            # Start navigation process
            process_key = f'navigation_{algorithm}'
            if process_key in self._launch_processes:
                # Kill existing process
                self._launch_processes[process_key].terminate()
                time.sleep(2)
            
            print(f"🚀 Starting navigation system: {' '.join(cmd)}")
            process = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            self._launch_processes[process_key] = process
            
            # Give it some time to start
            time.sleep(wait_time_s)  # Wait time for navigation system
            
            # Check if process is still running
            if process.poll() is None:
                return True, f"Navigation system started with {algorithm} algorithm", {
                    'algorithm': algorithm,
                    'process_id': process.pid,
                    'laser_bringup_pid': self._launch_processes[laser_bringup_key].pid,
                    'laser_odometry_pid': self._launch_processes[laser_odom_key].pid,
                    'map_path': map_path
                }
            else:
                return False, "Navigation system failed to start", {}
                
        except Exception as e:
            return False, f"Error starting navigation system: {str(e)}", {}
    
    def start_laser_bringup(self, wait_time_s: float = 10.0) -> tuple:
        """
        Start laser bringup separately
        
        Args:
            wait_time_s: Time to wait for initialization (default: 10.0 seconds)
        
        Returns:
            Tuple (success, message, data)
        """
        laser_bringup_key = 'laser_bringup'
        
        # Check if already running
        if laser_bringup_key in self._launch_processes and self._launch_processes[laser_bringup_key].poll() is None:
            return True, "Laser bringup already running", {
                'process_id': self._launch_processes[laser_bringup_key].pid
            }
        
        try:
            print("🔧 Starting laser bringup...")
            laser_bringup_cmd = ['ros2', 'launch', 'yahboomcar_nav', 'laser_bringup_launch.py']
            
            laser_bringup_process = subprocess.Popen(
                laser_bringup_cmd, 
                stdout=subprocess.DEVNULL, 
                stderr=subprocess.DEVNULL
            )
            self._launch_processes[laser_bringup_key] = laser_bringup_process
            
            # Wait for initialization
            time.sleep(wait_time_s)
            
            if laser_bringup_process.poll() is not None:
                return False, "Failed to start laser bringup. Make sure yahboomcar_nav laser_bringup_launch.py is available.", {}
            
            print("✅ Laser bringup started successfully")
            return True, "Laser bringup started successfully", {
                'process_id': laser_bringup_process.pid
            }
            
        except Exception as e:
            return False, f"Failed to start laser bringup: {str(e)}", {}
    
    def start_laser_odometry(self, wait_time_s: float = 5.0) -> tuple:
        """
        Start laser odometry separately (requires laser bringup to be running first)
        
        Args:
            wait_time_s: Time to wait for initialization (default: 5.0 seconds)
        
        Returns:
            Tuple (success, message, data)
        """
        laser_odom_key = 'laser_odometry'
        
        # Check if already running
        if laser_odom_key in self._launch_processes and self._launch_processes[laser_odom_key].poll() is None:
            return True, "Laser odometry already running", {
                'process_id': self._launch_processes[laser_odom_key].pid
            }
        
        # Check if laser bringup is running
        laser_bringup_key = 'laser_bringup'
        if laser_bringup_key not in self._launch_processes or self._launch_processes[laser_bringup_key].poll() is not None:
            return False, "Laser bringup must be started first. Please call start_laser_bringup() first.", {}
        
        try:
            print("🔧 Starting laser odometry...")
            laser_cmd = ['ros2', 'launch', 'rf2o_laser_odometry', 'rf2o_laser_odometry.launch.py']
            
            laser_process = subprocess.Popen(
                laser_cmd, 
                stdout=subprocess.DEVNULL, 
                stderr=subprocess.DEVNULL
            )
            self._launch_processes[laser_odom_key] = laser_process
            
            # Wait for initialization
            time.sleep(wait_time_s)
            
            if laser_process.poll() is not None:
                return False, "Failed to start laser odometry. Make sure rf2o_laser_odometry package is available.", {}
            
            print("✅ Laser odometry started successfully")
            return True, "Laser odometry started successfully", {
                'process_id': laser_process.pid
            }
            
        except Exception as e:
            return False, f"Failed to start laser odometry: {str(e)}", {}
    
    def stop_laser_bringup(self, force_kill: bool = False, wait_time_s: float = 3.0) -> tuple:
        """
        Stop laser bringup
        
        Args:
            force_kill: Whether to force kill the process
            wait_time_s: Time to wait before force killing (default: 3.0 seconds)
        
        Returns:
            Tuple (success, message, data)
        """
        laser_bringup_key = 'laser_bringup'
        
        if laser_bringup_key not in self._launch_processes:
            return False, "Laser bringup not running", {}
        
        try:
            process = self._launch_processes[laser_bringup_key]
            if process.poll() is None:  # Process is running
                if force_kill:
                    process.kill()
                else:
                    process.terminate()
                    time.sleep(wait_time_s)
                    
                    # Force kill if still running
                    if process.poll() is None:
                        process.kill()
                
                del self._launch_processes[laser_bringup_key]
                
                return True, "Laser bringup stopped successfully", {}
            else:
                del self._launch_processes[laser_bringup_key]
                return False, "Laser bringup was not running", {}
                
        except Exception as e:
            return False, f"Error stopping laser bringup: {str(e)}", {}
    
    def stop_laser_odometry(self, force_kill: bool = False, wait_time_s: float = 3.0) -> tuple:
        """
        Stop laser odometry
        
        Args:
            force_kill: Whether to force kill the process
            wait_time_s: Time to wait before force killing (default: 3.0 seconds)
        
        Returns:
            Tuple (success, message, data)
        """
        laser_odom_key = 'laser_odometry'
        
        if laser_odom_key not in self._launch_processes:
            return False, "Laser odometry not running", {}
        
        try:
            process = self._launch_processes[laser_odom_key]
            if process.poll() is None:  # Process is running
                if force_kill:
                    process.kill()
                else:
                    process.terminate()
                    time.sleep(wait_time_s)
                    
                    # Force kill if still running
                    if process.poll() is None:
                        process.kill()
                
                del self._launch_processes[laser_odom_key]
                
                return True, "Laser odometry stopped successfully", {}
            else:
                del self._launch_processes[laser_odom_key]
                return False, "Laser odometry was not running", {}
                
        except Exception as e:
            return False, f"Error stopping laser odometry: {str(e)}", {}
    
    def stop_navigation_system(self, include_laser_components: bool = True, force_kill: bool = False, wait_time_s: float = 5.0) -> tuple:
        """
        Stop the navigation system
        
        Args:
            include_laser_components: Whether to also stop laser bringup and odometry (default: True)
            force_kill: Whether to force kill processes
            wait_time_s: Time to wait before force killing (default: 5.0 seconds)
        
        Returns:
            Tuple (success, message, data)
        """
        try:
            stopped_processes = []
            
            # Stop navigation processes
            nav_processes = [key for key in self._launch_processes.keys() 
                           if key.startswith('navigation_')]
            
            for process_key in nav_processes:
                process = self._launch_processes[process_key]
                if process.poll() is None:  # Process is running
                    if force_kill:
                        process.kill()
                    else:
                        process.terminate()
                    stopped_processes.append(process_key)
            
            # Optionally stop laser components
            if include_laser_components:
                # Stop laser odometry first
                if 'laser_odometry' in self._launch_processes:
                    process = self._launch_processes['laser_odometry']
                    if process.poll() is None:
                        if force_kill:
                            process.kill()
                        else:
                            process.terminate()
                        stopped_processes.append('laser_odometry')
                
                # Then stop laser bringup
                if 'laser_bringup' in self._launch_processes:
                    process = self._launch_processes['laser_bringup']
                    if process.poll() is None:
                        if force_kill:
                            process.kill()
                        else:
                            process.terminate()
                        stopped_processes.append('laser_bringup')
            
            # Wait for processes to terminate (only if not force killing)
            if not force_kill:
                time.sleep(wait_time_s)
                
                # Force kill if still running
                for process_key in stopped_processes:
                    if process_key in self._launch_processes:
                        process = self._launch_processes[process_key]
                        if process.poll() is None:
                            process.kill()
            
            # Clean up process references
            for process_key in stopped_processes:
                if process_key in self._launch_processes:
                    del self._launch_processes[process_key]
            
            return True, f"Stopped processes: {stopped_processes}", {
                'stopped_processes': stopped_processes
            }
            
        except Exception as e:
            return False, f"Error stopping navigation system: {str(e)}", {}
    
    def __del__(self):
        """
        Cleanup when NavigationManager is destroyed
        """
        try:
            # Stop all processes
            self.stop_navigation_system(include_laser_components=True)
            
            # Cleanup ROS2 node
            if hasattr(self, '_ros_node') and self._ros_node:
                self._ros_node.destroy_node()
                
        except Exception as e:
            print(f"Error during NavigationManager cleanup: {e}")